"""
python fujian2_lstm_fillBlank.py
"""
import os
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime
from keras.models import Sequential
from keras.layers import LSTM, Dense, Input
from sklearn.preprocessing import MinMaxScaler

def lstm_fillBlank_single_category(input_path, output_path):
    
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    # 1. 加载数据
    with open(input_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # 将数据转为DataFrame格式
    df = pd.DataFrame(data)
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)

    # 2. 确定需要填补的日期范围
    fill_dates = pd.date_range(start='2022-10-01', end='2023-03-31')

    # 3. 准备数据
    known_data = df[(df.index >= '2022-07-01') & (df.index <= '2022-09-30')]
    known_data = pd.concat([known_data, df[(df.index >= '2023-04-01') & (df.index <= '2023-06-30')]])

    # 归一化数据
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(known_data[['sales']])

    # 创建训练集
    X, y = [], []
    for i in range(1, len(scaled_data)):
        X.append(scaled_data[i-1:i, 0])
        y.append(scaled_data[i, 0])
    X, y = np.array(X), np.array(y)
    X = X.reshape((X.shape[0], X.shape[1], 1))

    # 4. 构建LSTM模型
    model = Sequential()
    model.add(Input(shape=(X.shape[1], 1)))
    model.add(LSTM(100, activation='relu', return_sequences=True))  # 第一层LSTM，增加神经元并设置返回序列
    model.add(LSTM(50, activation='relu'))  # 第二层LSTM
    model.add(Dense(1))
    model.compile(optimizer='adam', loss='mse')

    # 5. 训练模型
    model.fit(X, y, epochs=500, verbose=0)

    # 6. 预测空白期的销量
    fill_predictions = []
    last_known = scaled_data[-1]
    for _ in fill_dates:
        last_known = model.predict(last_known.reshape((1, 1, 1)), verbose=0)
        fill_predictions.append(last_known[0][0])

    # 7. 反归一化
    predicted_sales = scaler.inverse_transform(np.array(fill_predictions).reshape(-1, 1))

    # 8. 保存填补后的数据
    filled_data = [{'date': date.strftime('%Y/%m/%d'), 'sales': int(sales)} 
                   for date, sales in zip(fill_dates, predicted_sales)]

    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(filled_data, f, ensure_ascii=False, indent=4)

    print(f"填补后的数据已保存至 {output_path}")

    # 9. 绘制图形
    plt.figure(figsize=(14, 7))

    # 绘制原始数据（散点图）
    plt.scatter(df.index, df['sales'], color='blue', s=10, label='原始数据')

    # 绘制填补的数据（散点图）
    fill_df = pd.DataFrame(filled_data)
    fill_df['date'] = pd.to_datetime(fill_df['date'])
    plt.scatter(fill_df['date'], fill_df['sales'], color='red', s=10, label='填补数据')

    # 图形美化
    plt.title('LSTM填补销量数据')
    plt.xlabel('日期')
    plt.ylabel('销量')
    plt.legend()
    plt.grid()
    plt.savefig(output_path.replace('.json', '_plot.png'))
    plt.show()
    print(f"图像已保存至 {output_path.replace('.json', '_plot.png')}")

if __name__ == '__main__':
    input_path = '../fujian/fujian2/groupByCategory/category_category1.json'
    output_path = '../fujian/fujian2/lstm_fillBlank/category_category1_filled.json'

    lstm_fillBlank_single_category(input_path, output_path)
